package com.maomao.es.vector;

import ai.djl.inference.Predictor;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelZoo;
import ai.djl.repository.zoo.ZooModel;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;

/**
 * @author mao
 * @date 2025-05-30- 14:57
 * @descript 可训练模型
 */
public class BertVectorizationService {

    private static final String MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2";
    private Criteria<String, float[]> criteria;
    private ZooModel<String, float[]> model;
    private Predictor<String, float[]> predictor;

    @PostConstruct
    public void init() throws Exception {
        criteria = Criteria.builder()
                .setTypes(String.class, float[].class)
                .optModelUrls("djl://ai.djl.huggingface.pytorch/" + MODEL_NAME)
                .build();
        model = ModelZoo.loadModel(criteria);
        predictor = model.newPredictor();
    }

    public float[] textToVector(String text) throws Exception {
        return predictor.predict(text);
    }

    @PreDestroy
    public void close() {
        if (predictor != null) {
            predictor.close();
        }
        if (model != null) {
            model.close();
        }
    }

}
